package com.cfp4cloud.cfp.knowledge.service.impl;

import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.cfp4cloud.cfp.knowledge.entity.AiDatasetEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiEmbedStoreEntity;
import com.cfp4cloud.cfp.knowledge.mapper.AiDatasetMapper;
import com.cfp4cloud.cfp.knowledge.mapper.AiEmbedStoreMapper;
import com.cfp4cloud.cfp.knowledge.service.EmbeddingStoreService;
import com.cfp4cloud.cfp.knowledge.support.handler.rag.EmbeddingStoreFactory;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * 嵌入向量存储服务实现类
 * <p>
 * 管理不同集合的向量存储实例 支持多种向量数据库类型 提供缓存机制优化访问性能
 *
 * @author chenda
 * @date 2024/9/26
 */
@Service
@RequiredArgsConstructor
public class EmbeddingStoreServiceImpl implements EmbeddingStoreService {

	/**
	 * 嵌入存储实例缓存 避免重复创建存储实例
	 */
	private final static Map<String, EmbeddingStore<TextSegment>> EMBEDDING_STORE_MAP = new HashMap<>();

	/**
	 * 嵌入存储工厂集合 按存储类型创建对应实例
	 */
	private final Map<String, EmbeddingStoreFactory> embeddingStoreFactoryMap;

	/**
	 * 嵌入存储配置数据访问层
	 */
	private final AiEmbedStoreMapper aiEmbedStoreMapper;

	/**
	 * 数据集数据访问层
	 */
	private final AiDatasetMapper aiDatasetMapper;

	/**
	 * 获取嵌入存储实例
	 * <p>
	 * 优先从缓存获取，缓存未命中则创建新实例 根据数据集配置选择对应的存储类型 无配置时默认使用内存存储
	 * @param collectionName 集合名称
	 * @return 嵌入存储实例
	 */
	@Override
	@SneakyThrows
	public EmbeddingStore<TextSegment> embeddingStore(String collectionName) {
		EmbeddingStore<TextSegment> embeddingStore = EMBEDDING_STORE_MAP.get(collectionName);

		if (Objects.nonNull(embeddingStore)) {
			return embeddingStore;
		}

		// 根据ID 查询知识库
		AiDatasetEntity aiDatasetEntity = aiDatasetMapper
			.selectOne(Wrappers.<AiDatasetEntity>lambdaQuery().eq(AiDatasetEntity::getCollectionName, collectionName));
		AiEmbedStoreEntity embedStoreEntity = aiEmbedStoreMapper.selectOne(Wrappers.<AiEmbedStoreEntity>lambdaQuery()
			.eq(AiEmbedStoreEntity::getStoreId, aiDatasetEntity.getStoreId()), false);

		// 没有配置信息，使用内存存储
		if (Objects.isNull(embedStoreEntity)) {
			embeddingStore = new InMemoryEmbeddingStore<>();
			EMBEDDING_STORE_MAP.put(collectionName, embeddingStore);
			return embeddingStore;
		}

		for (EmbeddingStoreFactory storeFactory : embeddingStoreFactoryMap.values()) {
			if (storeFactory.support(embedStoreEntity.getStoreType())) {
				embeddingStore = storeFactory.createEmbeddingStore(embedStoreEntity, aiDatasetEntity);
				break;
			}
		}

		EMBEDDING_STORE_MAP.put(collectionName, embeddingStore);
		return embeddingStore;
	}

	/**
	 * 删除嵌入存储缓存
	 * <p>
	 * 从缓存中移除指定集合的存储实例 下次访问时将重新创建
	 * @param name 集合名称
	 */
	@Override
	public void deleteEmbeddingStore(String name) {
		EMBEDDING_STORE_MAP.remove(name);
	}

	/**
	 * 批量删除向量数据
	 * <p>
	 * 从指定集合中删除多个向量 根据向量ID列表执行批量删除操作
	 * @param collectionName 集合名称
	 * @param idList 要删除的向量ID列表
	 */
	@Override
	@SneakyThrows
	public void delete(String collectionName, List<String> idList) {
		if (CollUtil.isEmpty(idList)) {
			return;
		}
		this.embeddingStore(collectionName).removeAll(idList);
	}

}
